import unittest

from tinydb.common.fabric import TableColumn
from tinydb.executor.operator.physical_operator import Sort, TableScan, HashAgg 

class TestSort(unittest.TestCase):
    def test_internal_sort(self):
        self._sort(sort_method=Sort.INTERNAL_SORT)

    def test_external_sort(self):
        self._sort(sort_method=Sort.EXTERNAL_SORT)

    def test_hash_agg(self):
        opt = HashAgg(
            group_by_column=TableColumn('t1', 'id'), 
            aggregate_function_name='count', 
            aggregate_column=TableColumn('t1', 'name'))
        opt.add_child(TableScan('t1'))
        # 到此位置，这个计划"树"就是：
        # Sort(group_column: id, agg_column: name, function: count)
        #   └── TableScan(table_name: t1)
        # 相当于执行了: select count(t1.name) from t1 group by t1.id
        opt.open()
        results = list(opt.next())
        expected = [(1, 1), (2, 1), (3, 1), (4, 1)] # id=1,2,3,4 的记录各一条，所以 count=1
        self.assertEqual(results, expected)
        opt.close()
        

    def _sort(self, sort_method):
        # name 排序(字符串排序)
        opt = Sort(TableColumn('t1', 'name'))
        opt.method = sort_method
        opt.add_child(TableScan('t1'))
        # 到此位置，这个计划"树"就是：
        # Sort(sort_column: name)
        #   └── TableScan(table_name: t1)
        opt.open()
        results = list(opt.next())
        expected = [(4, '小丽'), (3, '小刚'), (1, '小明'), (2, '小红')]
        self.assertEqual(results, expected)
        opt.close()

        opt = Sort(TableColumn('t1', 'name'), asc=False)
        opt.method = sort_method
        opt.add_child(TableScan('t1'))
        # 到此位置，这个计划"树"就是：
        # Sort(sort_column: name)
        #   └── TableScan(table_name: t1)
        opt.open()
        results = list(opt.next())
        expected = [(2, '小红'), (1, '小明'), (3, '小刚'), (4, '小丽')]
        self.assertEqual(results, expected)
        opt.close()

        # id 排序
        opt = Sort(TableColumn('t1', 'id'), asc=False)
        opt.method = sort_method
        opt.add_child(TableScan('t1'))
        # 到此位置，这个计划"树"就是：
        # Sort(sort_column: id)
        #   └── TableScan(table_name: t1)
        opt.open()
        results = list(opt.next())
        
        expected = [(4, '小丽'), (3, '小刚'), (2, '小红'), (1, '小明')]
        self.assertEqual(results, expected)
        opt.close()

        opt = Sort(TableColumn('t1', 'id'), asc=True)
        opt.method = sort_method
        opt.add_child(TableScan('t1'))
        # 到此位置，这个计划"树"就是：
        # Sort(sort_column: id)
        #   └── TableScan(table_name: t1)
        opt.open()
        results = list(opt.next())
        
        expected = [(1, '小明'), (2, '小红'), (3, '小刚'), (4, '小丽')]
        self.assertEqual(results, expected)
        opt.close()

if __name__ == '__main__':
    unittest.main()
